"""
使用已训练好的 KNN 模型构建一个 Gradio Web 应用：
- 允许用户在手写板上写 0~9 的数字；
- 将输入预处理为与 sklearn digits(8x8) 相同的格式并做强度归一；
- 给出预测结果。

兼容性说明：不同版本的 gradio 对手写板组件的实现可能不同。这里优先使用
`gr.Sketchpad`（高版本），若不可用则回退到 `gr.Image`。
"""

from __future__ import annotations

import io
import base64
import pickle
from pathlib import Path
from typing import Optional

import gradio as gr
import numpy as np
from PIL import Image


ROOT = Path(__file__).resolve().parent
MODEL_PATH = ROOT / "best_knn_model.pkl"


def load_model():
	if not MODEL_PATH.exists():
		raise FileNotFoundError(
			f"未找到模型文件: {MODEL_PATH}. 请先运行 optimal_knn.py 训练并保存模型。"
		)
	with open(MODEL_PATH, "rb") as f:
		return pickle.load(f)


model = None
try:
	model = load_model()
except Exception as e:
	# 延迟到首次预测时再报错，便于 UI 启动
	model = None
	load_error = str(e)


def preprocess_to_digits_like(img: Image.Image) -> np.ndarray:
	"""将任意图像转换为 sklearn digits 样式的 8x8 灰度向量(长度64)。"""
	# 转灰度
	img = img.convert("L")
	# 反色：白底黑字 -> 黑底白字（sklearn digits 更接近亮笔画）
	np_img = np.array(img).astype(np.float32)
	# 将画布背景当作白色(255)，前景为较小值，先反转
	np_img = 255.0 - np_img
	# 缩放到 8x8
	img_small = Image.fromarray(np.clip(np_img, 0, 255).astype(np.uint8)).resize((8, 8), Image.BILINEAR)
	arr = np.array(img_small).astype(np.float32)
	# 归一到 0-16，digits 数据特征范围大致为 0..16
	arr = arr / 255.0 * 16.0
	# 展平为 (64,)
	return arr.reshape(1, -1)


def predict_fn(image: Optional[np.ndarray | Image.Image]):
	"""预测函数：始终返回整数；异常使用 gr.Error 抛出，避免向 Number 返回字符串。"""
	# 输入可能来自 Sketchpad（numpy RGBA/L/float）或 Image（PIL Image）
	if image is None:
		raise gr.Error("请先在画布上写一个数字。")

	# 延迟加载失败时给出清晰错误
	global model
	if model is None:
		try:
			model = load_model()
		except Exception as e:  # pragma: no cover - 仅运行时提示
			raise gr.Error(f"模型加载失败: {e}")

	# 将输入统一为 PIL Image（兼容 dict/list/ndarray、2D/3D、浮点/整型、含 alpha）
	if isinstance(image, dict):
		# 兼容不同 gradio 版本下的返回结构
		candidate = None
		for key in ("image", "composite", "background"):
			if key in image and image[key] is not None:
				candidate = image[key]
				break
		if candidate is None and "layers" in image and image["layers"]:
			# 图层合成：将每个 RGBA 层叠加到白色背景
			layers = image["layers"]
			# 统一为 numpy
			norm_layers = []
			for lyr in layers:
				if isinstance(lyr, list):
					lyr = np.array(lyr)
				norm_layers.append(np.asarray(lyr))
			# 取最大尺寸
			h = max(l.shape[0] for l in norm_layers)
			w = max(l.shape[1] for l in norm_layers)
			canvas = Image.new("RGBA", (w, h), (255, 255, 255, 255))
			for lyr in norm_layers:
				arr = lyr
				if arr.dtype != np.uint8:
					maxv = float(np.max(arr)) if arr.size else 1.0
					if maxv <= 1.0:
						arr = (arr * 255.0).clip(0, 255).astype(np.uint8)
					else:
						arr = np.clip(arr, 0, 255).astype(np.uint8)
				if arr.ndim == 2:
					pil = Image.fromarray(arr, mode="L").convert("RGBA")
				elif arr.ndim == 3:
					if arr.shape[-1] == 4:
						pil = Image.fromarray(arr, mode="RGBA")
					else:
						pil = Image.fromarray(arr[:, :, :3], mode="RGB").convert("RGBA")
				else:
					continue
				canvas.alpha_composite(pil)
			image = canvas.convert("RGB")
		else:
			# candidate 可能是 list/ndarray/base64/pil
			if isinstance(candidate, list):
				candidate = np.array(candidate)
			if isinstance(candidate, (bytes, bytearray)):
				try:
					image = Image.open(io.BytesIO(candidate)).convert("RGB")
				except Exception:
					raise gr.Error("无法解析图像字节流。")
			elif isinstance(candidate, str) and candidate.startswith("data:image"):
				# data URL base64
				try:
					b64 = candidate.split(",", 1)[1]
					image = Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")
				except Exception:
					raise gr.Error("无法解析图像Base64。")
			elif isinstance(candidate, np.ndarray):
				image = candidate
			elif isinstance(candidate, Image.Image):
				image = candidate
			else:
				# 回退为原对象尝试后续分支
				image = candidate

	if isinstance(image, list):
		image = np.array(image)

	if isinstance(image, np.ndarray):
		arr = image
		if arr.dtype != np.uint8:
			# 归一到 0..255 再转 uint8
			maxv = float(np.max(arr)) if arr.size else 1.0
			if maxv <= 1.0:
				arr = (arr * 255.0).clip(0, 255).astype(np.uint8)
			else:
				arr = np.clip(arr, 0, 255).astype(np.uint8)

		if arr.ndim == 2:
			pil = Image.fromarray(arr, mode="L")
		elif arr.ndim == 3:
			if arr.shape[-1] == 4:
				pil = Image.fromarray(arr, mode="RGBA")
				bg = Image.new("RGBA", pil.size, (255, 255, 255, 255))
				pil = Image.alpha_composite(bg, pil).convert("RGB")
			else:
				pil = Image.fromarray(arr[:, :, :3], mode="RGB")
		else:
			raise gr.Error("无法识别的图像维度。")
		image = pil
	elif not isinstance(image, Image.Image):
		raise gr.Error("无法识别的图像格式。")

	feat = preprocess_to_digits_like(image)
	try:
		pred = model.predict(feat)[0]
	except Exception as e:
		raise gr.Error(f"预测失败: {e}")
	return int(pred)


with gr.Blocks(title="KNN 手写数字识别") as demo:
	gr.Markdown("# 使用 KNN 的手写数字识别\n在画布上写一个 0~9 的数字后点击提交。")

	with gr.Row():
		try:
			sketch = gr.Sketchpad(label="手写板", image_mode="L", brush=20, height=256, width=256)
			input_comp = sketch
		except Exception:
			# 回退方案：使用 Image 作为输入
			img = gr.Image(label="手写板(兼容模式)", type="pil", tool=None, image_mode="L", height=256, width=256)
			input_comp = img

		out = gr.Number(label="预测结果", precision=0)

	submit = gr.Button("提交")
	submit.click(fn=predict_fn, inputs=input_comp, outputs=out)

	# 如果模型提前加载失败，给出提示
	if 'load_error' in globals():
		gr.Markdown(f"> ⚠️ {load_error}")


if __name__ == "__main__":
	# 默认本地运行；如需公网分享可设置 share=True
	demo.launch()
